## ----------------------------- a naive guide for W matrix   -------------------------------------------
## 1. check sparisity: row count and column count of neighbors with positive weights
## 2. check eigen value: the minimum negative eigenvalue should be around -1, 
##                       if the minimum negative eigenvalue is around 0, then the W matrix may not work
## 3. w matrix with complex eigen value (asymmetric w matrix): no worries, just 
##                                                             check 2 for all real eigenvalues
##-------------------------------------------------------------------------------------------------------


## ----------------------The MCMC output includes posteriors of the following parameters  --------------

## [1] "beta": the invariant part of coefficients
## [2]"rhoy": autoregressive coefficient associated to y lag
## [3] "omega": variance parameters of factor loadings for factor selection
## [4] "sigma2": error variance
## [5] "Alpha": unit-varying part of coeficients  
## [6] "Xi": time-varying part of coeficients
## [7] "rho0": constant sar coef
## [8] "Rho": random sar coef
## [9] "kappa": ar1 parameter for rho_t (when using AR(1) precess)
## [10] "sigma_n2": error variance in state space rho_t
## [11] "rhoA":  coef for time-vary covar that explains rho_t
## [12] "L": factor loadings matrix 
## [13] "Factor":  factor matrix 
##------------------------------------------------------------------------------------------------------


## ---------------------- sim1 block diagnoal w matrix
rm(list = ls())

sink("log/log_block_sim3.txt") # start log
begin.time<-Sys.time()


## compile functions 
library(bpNet)
##library(mcmcr)
library(coda)
library(MASS)
library(abind)

## Set working directory where the data, code, and output are stored
set.seed(123456789)

## ------------------------- ##
source("code/simulated20529.R")   ## code for generating simulated data
source("code/net_plot.R")

## 1. The True Values of parameters except rho_t 

   N <- 200 ## number of units
   Ng <- 20 ## number of groups 
   TT <- 50 ## time period
   unit <- N / Ng

   p <- 2
   beta <- as.matrix(c(5, 1, 3, -2, 2))

## rho_t
   rhoA <- as.matrix(c(0.3, -0.2))
   kappa <- 0.3
   sigma_n2 <- 0.36

## factor numbers
   r <- 2

## error variance
   sigma2 <- 1

## Load W
    load('data/SimData/blockW.RData')
## Load Simulation Data
    load('data/SimData/sim1cor3.RData')
## get the data for simulation I
    data1 <- sim[[1]]$data
    ## true rho (for comparison)
    Rho1 <- sim[[1]]$Rho
    load('data/SimData/sim2cor3.RData')
    ## get the data for simulation II
    data2 <- sim2[[1]]$data

    ## true rho
    Rho2 <- sim2[[1]]$Rho
    ## covariates in the rho_t process
    rhoZ <- sim2[[2]]

    ## number of iterations (mcmc) after the burnin stage (burnin)
    ##mcmc <- 25000
    ##burnin <- 10000
    mcmc <- 15000
    burnin <- 5000 

    ## ------------ Simulation Study I: Comparing Four Models
    ## M1: with 2 factors and no Bayesian shrinkage (True Model)
    out1.1 <- bpNet(data = data1,   ## a data frame
	              W = W,          ## an array N * N * TT, each slide is a w matrix for a given period
	              index = c("id", "time"), ## id and time
	              Yname = "y",
	              Xname = c("x1", "x2"),
	              Zname = NULL,   ## covar has unit random effect
	              Aname = NULL,   ##           time
	              Contextual = c("x1", "x2"),  ## contectual covar
	              Contextual.effect = "both",  ## contectual effect: "none", "unit", "time", "both"
	              lagY = FALSE,                ## whether to incorporate lagged outcome
	              force = "both",              ## two-way fe" "none", "unit", "time", "both"
	              r = 2,                       ## number of factors
	              flasso = 0,                   ## with/without factor selection (1/0)
	              rhoZ = as.matrix(rhoZ),       ## time-varying covar that explains rho: T*k    
	              rho_spec = "ar1",                ## options of rho_t process: "multilevel", "rw", "ar1"
	              niter = mcmc,
	              burn = burnin)

     ## M2: with 10 factors and Bayesian shrinkage
     out1.2 <- bpNet(data = data1,   
	              W = W,         
	              index = c("id", "time"), 
	              Yname = "y",
	              Xname = c("x1", "x2"),
	              Zname = NULL,  
	              Aname = NULL,   
	              Contextual = c("x1", "x2"), 
	              Contextual.effect = "both", 
	              lagY = FALSE,               
	              force = "both",            
	              r = 10,                    ## number of factors
	              flasso = 1,                ## with/without factor selection (1/0)
	              rhoZ = as.matrix(rhoZ),     
	              rho_spec = "ar1",               
	              niter = mcmc,
	              burn = burnin)

     ## M3: with two-way fixed effects but no factors
     out1.3 <- bpNet(data = data1,  
	              W = W,        
	              index = c("id", "time"), 
	              Yname = "y",
	              Xname = c("x1", "x2"),
	              Zname = NULL,   
	              Aname = NULL,  
	              Contextual = c("x1", "x2"), 
	              Contextual.effect = "both", 
	              lagY = FALSE,               
	              force = "both",             
	              r = 0,                   ## number of factors
	              flasso = 0,              ## with/without factor selection (1/0)
	              rhoZ = as.matrix(rhoZ),       
	              rho_spec = "ar1",           
	              niter = mcmc,
	              burn = burnin)



    ## M4: constant rho, 2 factors, and no shrinkage
    out1.4 <- bpNet(data = data1,  
	              W = W,         
	              index = c("id", "time"), 
	              Yname = "y",
	              Xname = c("x1", "x2"),
	              Zname = NULL,   
	              Aname = NULL,   
	              Contextual = c("x1", "x2"),  
	              Contextual.effect = "both",  
	              lagY = FALSE,                
	              force = "both",             
	              r = 2,
	              flasso = 0,             
	              rhoZ = as.matrix(rhoZ),      
	              rho_spec = "constant",               
	              niter = mcmc,
	              burn = burnin)

   ## ---------------Plot MCMC results------------ @@
    ### Plot rho or rho_t
    ## --------------- plot rho_t
    ##  Figure 2 in the main text: Study I
    a <- 25  # set the font
    p1.1 <- rho_plot(out1.1, Rho1, constant = FALSE)
    p1.1 <- p1.1 +theme(axis.text.x = element_text( hjust = 1, size=a))+ xlab(" ") + ylab(" ")
    p1.1 <- p1.1 +theme(axis.text.y = element_text( hjust = 1, size=a))
    p1.1 <- p1.1 +theme(legend.title = element_text(size=a), legend.text = element_text(size=a), legend.key.size = unit(2, 'cm'))

    p1.2 <- rho_plot(out1.2, Rho1, constant = FALSE)
    p1.2 <- p1.2 +theme(axis.text.x = element_text( hjust = 1, size=a))+ xlab(" ") + ylab(" ")
    p1.2 <-p1.2 +theme(axis.text.y = element_text( hjust = 1, size=a))
    p1.2 <-p1.2 +theme(legend.title = element_text(size=a), legend.text = element_text(size=a), legend.key.size = unit(2, 'cm'))
    
   p1.3 <- rho_plot(out1.3, Rho1, constant = FALSE)
   p1.3 <- p1.3 +theme(axis.text.x = element_text( hjust = 1, size=a))+ xlab(" ") + ylab(" ")
   p1.3 <- p1.3 +theme(axis.text.y = element_text( hjust = 1, size=a))
   p1.3 <- p1.3 +theme(legend.title = element_text(size=a), legend.text = element_text(size=a), legend.key.size = unit(2, 'cm'))
    
   p1.4 <- rho_plot(out1.4, Rho1, constant = TRUE)
   p1.4 <- p1.4 +theme(axis.text.x = element_text( hjust = 1, size=a))+ xlab(" ") + ylab(" ")
   p1.4 <- p1.4 +theme(axis.text.y = element_text( hjust = 1, size=a))
   p1.4 <- p1.4+theme(legend.title = element_text(size=a), legend.text = element_text(size=a), legend.key.size = unit(2, 'cm'))
    ggsave('plots/sim1cor3/rho1.pdf', p1.1, width = 10, height = 6)
    ggsave('plots/sim1cor3/rho2.pdf', p1.2, width = 10, height = 6)
    ggsave('plots/sim1cor3/rho3.pdf', p1.3, width = 10, height = 6)
    ggsave('plots/sim1cor3/rho4.pdf', p1.4, width = 10, height = 6)

    ## may use hist of the constant rho
    #pdf('plots/sim1cor3/rho4hist.pdf',width=18, height=9, pointsize = 20)
    #rhoC <- out1.4$rho0
    #CC <- as.mcmc(t(rhoC))
    #hist(CC, col="gray90", main=" ",  xlab=expression(rho), cex.axis=1.5)
    #dev.off()

      ## to release memory and remove the mcmc output in Study I
      remove(out1.1, out1.2, out1.3, out1.4)

##############################################################################
       ## -------------Study II: rho=0 ------------------
      ## M1: 2 factors without shrinkage
      out2.1 <- bpNet(data = data2,   
	              W = W,          
	              index = c("id", "time"), 
	              Yname = "y",
	              Xname = c("x1", "x2"),
	              Zname = NULL,   
	              Aname = NULL,   
	              Contextual = c("x1", "x2"),  
	              Contextual.effect = "both",  
	              lagY = FALSE,               
	              force = "both",             
	              r = 2,
	              flasso = 0,              
	              rhoZ = as.matrix(rhoZ),        
	              rho_spec = "ar1",                
	              niter =mcmc,
	              burn = burnin)


      ## M2: 10 factors and Bayesian shrinkage
      out2.2 <- bpNet(data = data2,   
	              W = W,          
	              index = c("id", "time"), 
	              Yname = "y",
	              Xname = c("x1", "x2"),
	              Zname = NULL,   
	              Aname = NULL,   
	              Contextual = c("x1", "x2"),  
	              Contextual.effect = "both",  
	              lagY = FALSE,                
	              force = "both",              
	              r = 10,
	              flasso = 1,              
	              rhoZ = as.matrix(rhoZ),       
	              rho_spec = "ar1",               
	              niter = mcmc,
	              burn = burnin)



     ## M3: two-way fixed effect, no factors
     out2.3 <- bpNet(data = data2,  
	              W = W,          
	              index = c("id", "time"), 
	              Yname = "y",
	              Xname = c("x1", "x2"),
	              Zname = NULL,   
	              Aname = NULL,   
	              Contextual = c("x1", "x2"),  
	              Contextual.effect = "both",  
	              lagY = FALSE,               
	              force = "both",             
	              r = 0,
	              flasso = 0,             
	              rhoZ = as.matrix(rhoZ),      
	              rho_spec = "ar1",               
	              niter = mcmc,
	              burn = burnin)

     ## M4: constant rho, two factors without shrinkage
     out2.4 <- bpNet(data = data2,  
	              W = W,         
	              index = c("id", "time"), 
	              Yname = "y",
	              Xname = c("x1", "x2"),
	              Zname = NULL,  
	              Aname = NULL,   
	              Contextual = c("x1", "x2"), 
	              Contextual.effect = "both",  
	              lagY = FALSE,                
	              force = "both",              
	              r = 2,
	              flasso = 0,             
	              rhoZ = as.matrix(rhoZ),      
	              rho_spec = "constant",              
	              niter = mcmc,
	              burn = burnin)

    ## ---------------Plot MCMC results------------ @@

    ##  Figure 4 in the main text: Study II
    p2.1 <- rho_plot(out2.1, Rho2, constant = FALSE)
    p2.1 <- p2.1 +theme(axis.text.x = element_text( hjust = 1, size=a))+ xlab(" ") + ylab(" ")
    p2.1 <- p2.1 +theme(axis.text.y = element_text( hjust = 1, size=a))
    p2.1 <- p2.1 +theme(legend.title = element_text(size=a), legend.text = element_text(size=a), legend.key.size = unit(2, 'cm'))

    p2.2 <- rho_plot(out2.2, Rho2, constant = FALSE)
    p2.2 <- p2.2 +theme(axis.text.x = element_text( hjust = 1, size=a))+ xlab(" ") + ylab(" ")
    p2.2 <- p2.2 +theme(axis.text.y = element_text( hjust = 1, size=a))
    p2.1 <- p2.1 +theme(legend.title = element_text(size=a), legend.text = element_text(size=a), legend.key.size = unit(2, 'cm'))

    p2.3 <- rho_plot(out2.3, Rho2, constant = FALSE )
    p2.3 <- p2.3 +theme(axis.text.x = element_text( hjust = 1, size=a))+ xlab(" ") + ylab(" ")
    p2.3 <- p2.3 +theme(axis.text.y = element_text( hjust = 1, size=a))
    p2.1 <- p2.1 +theme(legend.title = element_text(size=a), legend.text = element_text(size=a), legend.key.size = unit(2, 'cm'))

    p2.4 <- rho_plot(out2.4, Rho2, constant = TRUE)
    p2.4 <- p2.4 +theme(axis.text.x = element_text( hjust = 1, size=a))+ xlab(" ") + ylab(" ")
    p2.4 <- p2.4 +theme(axis.text.y = element_text( hjust = 1, size=a))
    p2.1 <- p2.1 +theme(legend.title = element_text(size=a), legend.text = element_text(size=a), legend.key.size = unit(2, 'cm'))

    ggsave('plots/sim2cor3/rho5.pdf', p2.1, width = 10, height = 6)
    ggsave('plots/sim2cor3/rho6.pdf', p2.2, width = 10, height = 6)
    ggsave('plots/sim2cor3/rho7.pdf', p2.3, width = 10, height = 6)
    ggsave('plots/sim2cor3/rho8.pdf', p2.4, width = 10, height = 6)
    
    ### Plot the constant rho
    pdf('plots/sim2cor3/rho8hist.pdf',width=18, height=9, pointsize = 20)
    rhoC <- out2.4$rho0
    CC <- as.mcmc(t(rhoC))
    hist(CC, col="gray90", main=" ",  xlab=expression(rho),  cex.axis=1.5)
    abline(v=0, col = "red", lwd=6, lty=2)
    dev.off()

print(Sys.time()-begin.time) 
sink() # close log    

